% Auxiliary Function: hybrid_fmincon2
%   Description:
%       This function performs a two-stage optimization procedure to minimize the objective function
%       with nonlinear constraints, where it uses a global method (e.g., pswarm) in the first stage
%       and local refinement (fmincon) in the second stage
%   Inputs:
%       crit_fn: Objective function to minimize
%       x0s: Initial value(s) for optimization, which can be selected by first running pswarm
%       lb: Lower bounds for variables
%       ub: Upper bounds for variables
%       opt_opts: Struct with options like maximum interation, used methods, etc.
%       nl_constr_fun: Nonlinear constraint function
%   Outputs:
%       x: Solution vector (minimizer)
%       fval: Objective function value at x
%       exitflag: Indicates the convergence status
%       output: Struct with additional information, e.g., iterations, etc.

function [x, fval, exitflag, output] = hybrid_fmincon2(crit_fn, x0s, lb, ub, opt_opts, nl_constr_fun)
    % ========================================
    % Parameter setting
    % ========================================
    % Determine the number of variables to optimize
    dim_x = size(x0s,1);
    % Check if nonlinear constraint is imposed
    if nargin<6
        nl_constr_fun = [];
    else
        if ~isempty(nl_constr_fun)
            assert(isa(nl_constr_fun, 'function_handle'));
        end
    end
    % Other settings
    display_str = get_default_field(opt_opts, 'Display', 'off');
    tic_ID = tic;
    if ~isempty(lb)
        lb = lb(:);
    end
    if ~isempty(ub)
        ub = ub(:);
    end
    % ========================================
    % Step 1: Global method
    % ========================================
    % ==============================
    % Use genetic algorithm (GA)
    % ==============================
    % Comments on GA:
    %   A problem when there are nonlinear constraints is that GA population does not always satisfy
    %   them, despite the initial population does
    %   If we do not run GA until convergence, the "fittest" value may not be ideal
    if strcmpi(opt_opts.method1_name, 'GA')
        ga_PopulationSize = get_default_field(opt_opts, 'ga_PopulationSize_step1', max(50, min(200, dim_x * 10)));
        % Generate initial population manually, as the default in GA is rather wild
        n_ga_generate_addl_Population = ga_PopulationSize - size(x0s, 2);
        if n_ga_generate_addl_Population > 0 
            if isempty(nl_constr_fun)
                ga_x0s = [x0s, make_b0_lb_ub(lb, ub, 0, n_ga_generate_addl_Population)];
            else
                % GA does not satisfy nonlinear constraints in general
                % Therefore, we try to generate good initial population
                % The code below assumes that we only have inequality constraints
                % Try many values as we assume that the cost of computing nl_constr_fun is small
                n_try = 100 + 10 * n_ga_generate_addl_Population;
                try_x0s = make_b0_lb_ub(lb, ub, 0, n_try);
                ix_ok = false(n_try, 1);
                n_ok = 0;
                for i_x = 1:n_try
                    nlc = nl_constr_fun(try_x0s(:, i_x));
                    if max(nlc) <= 0
                        ix_ok(i_x) = 1;
                        n_ok = n_ok + 1;
                        if n_ok >= n_ga_generate_addl_Population
                            break;
                        end
                    end
                end
                % After the procedure, the number of column may be smaller than ga_PopulationSize,
                % so GA will generate the rest
                ga_x0s = [x0s, try_x0s(:, ix_ok)]; 
            end
        end
        opts_ga = optimoptions('ga', ...
            'InitialPopulation', ga_x0s', ...
            'Display', display_str, ...
            'PopulationSize', ga_PopulationSize, ...
            'ConstraintTolerance', 1e-6, ...
            'MaxGenerations', get_default_field(opt_opts, 'MaxIter_step1', 100));
        [step1.x, step1.fval, step1.exitflag, step1.output] = ga(crit_fn, dim_x, [], [], [], [], ...
            lb, ub, nl_constr_fun, opts_ga);
    % ==============================
    % Use pswarm
    % ==============================
    elseif strcmpi(opt_opts.method1_name, 'PSwarm1')
        assert(isempty(nl_constr_fun));
        % Define optimization problem in the PSwarm1 format
        Problem.ObjFunction = crit_fn;  
        Problem.LB = lb; 
        Problem.UB = ub; 
        % Set up initial values
        for j = size(x0s, 2):-1:1 
            InitPop(j).x = x0s(:, j);
        end
        % Extract optimization parameters
        Options.Size    = get_default_field(opt_opts, 'pswarm_size', min(42, 10 + 5 * dim_x));
        Options.MaxIter = get_default_field(opt_opts, 'MaxIter_step1', 100);
        Options.MaxObj  = get_default_field(opt_opts, 'MaxFunEvals_step1', 1000 * (1 + dim_x));
        if strcmp(display_str, 'iter')
            Options.IPrint = 10;
        elseif strcmp(display_str, 'off')
            Options.IPrint = -1;
        else
            Options.IPrint = 0;
        end
        % Run PSwarm1
        [step1.x, step1.fval, step1.output] = PSwarm1(Problem, InitPop, Options); 
        step1.exitflag = 0;
    % ==============================
    % Skip step 1 and just uses all given x0s in step 2 directly
    % ==============================
    % Multistart fmincon can be implemented here
    elseif strcmpi(opt_opts.method1_name, 'multi_fmincon')
        step1.x = x0s;
        step1.fval = 1e100;
    % ==============================
    % Pop error if an unknown method is supplied
    % ==============================
    else
        error('Unknown method1_name = "%s"', opt_opts.method1_name);
        assert(isvector(x0s));
        step1.x = x0s;
        step1.fval = crit_fn(step1.x);
    end
    % Post step 1 processing
    if strcmpi(opt_opts.method1_name, 'multi_fmincon') 
        x0s_fmincon = x0s;
    else
        assert(isvector(step1.x));
        if get_default_field(opt_opts, 'step2_multi_fmincon', 1)
            x0s_fmincon = [x0s, step1.x(:)];
        else
            x0s_fmincon = step1.x(:);
        end
    end
    step1.took = toc(tic_ID);
    % ========================================
    % Step 2: Local refinement with fmincon
    % ========================================
    % Set up parameters for fmincon
    fmincon_opts = optimset( ...
        'Display', display_str, ...
        'Algorithm',   get_default_field(opt_opts, 'Algorithm', 'sqp'), ...
        'MaxFunEvals', get_default_field(opt_opts, 'MaxFunEvals_fmincon', []), ...
        'MaxIter',     get_default_field(opt_opts, 'MaxIter_fmincon', []), ...
        'TolX',        get_default_field(opt_opts, 'TolX', 1e-6), ...
        'TolFun',      get_default_field(opt_opts, 'TolFun', 1e-6));
    % Create empty arrays for output
    best_fval = 1e100;
    best_x = [];
    best_exitflag = []; 
    best_output = [];
    % Optimization procedure with fmincon
    for j_sv = 1:size(x0s_fmincon, 2)
        [x, fval, exitflag, output] = fmincon( ...
            crit_fn, x0s_fmincon(:, j_sv), [], [], [], [], ...
            lb, ub, nl_constr_fun, fmincon_opts);
        if fval < best_fval
            best_x = x;
            best_fval = fval;
            best_exitflag = exitflag;
            best_output = output;
        end
    end
    % ========================================
    % Prepare result for output
    % ======================================== 
    x = best_x;
    fval = best_fval;
    exitflag = best_exitflag;
    output = best_output;
    if isempty(nl_constr_fun)
        assert(fval <= step1.fval + 1e-7);
    elseif fval > step1.fval + 1e-7
        warning('fval = %5.5g > step1.fval = %5.5g', fval, step1.fval);
    end
    output.step1 = step1;
    output.took_total = toc(tic_ID);
end

